from utils.tools import session, logger
import requests
import json
import os
from feapder.utils.tools import get_file_list
from config import (
    conf,
    qanything_api_list_knowledge_base,
    qanything_api_new_knowledge_base,
    qanything_api_upload_files,
    qanything_api_list_files,
    qanything_api_get_bot_info,
    qanything_api_new_bot,
    qanything_headers,
)

"""
qa bot init

获取文件列表
获取知识库

创建知识库
上传文件
创建机器人
机器人信息查看
"""


class KnowledgeBaseClient:
    def __init__(self):
        self.url = qanything_api_list_knowledge_base
        self.headers = qanything_headers

    def _prepare_request_data(self):
        """准备请求数据"""
        return {"user_id": conf.get_value("QANYTHING_USER_ID")}

    def _send_request(self, url, headers, data):
        """发送请求"""
        try:
            response = session.post(url, headers=headers, data=json.dumps(data))
            return response
        except requests.RequestException as e:
            logger.error(f"请求错误: {e}")
            raise

    def _handle_response(self, response):
        """处理响应"""
        try:
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            logger.error(f"HTTP错误: {http_err}")
        except Exception as err:
            logger.error(f"其他错误: {err}")

    def get_knowledge_base(self):
        """获取知识库列表"""
        data = self._prepare_request_data()
        try:
            response = self._send_request(self.url, self.headers, data)
            return self._handle_response(response)
        except requests.RequestException as e:
            logger.error(f"请求失败: {e}")


class NewKnowledgeBaseClient:

    def __init__(self, kb_name) -> None:
        self.url = qanything_api_new_knowledge_base
        self.headers = qanything_headers
        self.user_id = conf.get_value("QANYTHING_USER_ID")
        self.kb_name = kb_name

    def _prepare_request_data(self):
        """准备请求数据"""
        return {"user_id": self.user_id, "kb_name": self.kb_name}

    def _send_request(self, url, headers, data):
        """发送请求"""
        try:
            response = session.post(url, headers=headers, data=json.dumps(data))
            return response
        except requests.RequestException as e:
            logger.error(f"请求错误: {e}")
            raise

    def _handle_response(self, response):
        """处理响应"""
        try:
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            logger.error(f"HTTP错误: {http_err}")
        except Exception as err:
            logger.error(f"其他错误: {err}")

    def new_knowledge_base(self):
        """新建知识库"""
        data = self._prepare_request_data()
        try:
            response = self._send_request(self.url, self.headers, data)
            return self._handle_response(response)
        except requests.RequestException as e:
            logger.error(f"请求失败: {e}")


class KnowledgeBaseUploadClient:

    def __init__(self, folder_path, kb_id) -> None:
        self.url = qanything_api_upload_files
        self.user_id = conf.get_value("QANYTHING_USER_ID")
        self.folder_path = folder_path
        self.kb_id = kb_id

    def _prepare_request_data(self):
        """准备请求数据"""
        """
        上传模式
        1.soft:知识库内存在同名文件时当前文件不再上传
        2.strong:文件名重复的文件强制上传
        """

        return {"user_id": self.user_id, "kb_id": self.kb_id, "mode": "soft"}

    def _send_request(self, url, files, data):
        """发送请求"""
        try:
            response = session.post(url, files=files, data=data)
            return response
        except requests.RequestException as e:
            logger.error(f"请求错误: {e}")
            raise

    def _handle_response(self, response):
        """处理响应"""
        try:
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            logger.error(f"HTTP错误: {http_err}")
        except Exception as err:
            logger.error(f"其他错误: {err}")

    def _find_files(self):
        files = []
        print("self.folder_path", self.folder_path)
        for file_path in get_file_list(self.folder_path):
            if file_path.endswith(
                ".txt"
            ):  # 这里只上传后缀是txt的文件，请按需修改，支持类型：
                files.append(("files", open(file_path, "rb")))

        if len(files) == 0:
            raise Exception("没有找到文件")
        logger.info(f"读取到{len(files)}个文件")
        return files

    def upload_files(self):
        """上传文件"""
        files = self._find_files()
        data = self._prepare_request_data()
        try:
            response = self._send_request(self.url, files, data)
            return self._handle_response(response)
        except requests.RequestException as e:
            logger.error(f"请求失败: {e}")


class KnowledgeBaseListFiles:

    def __init__(self, kb_id) -> None:
        self.url = qanything_api_list_files
        self.user_id = conf.get_value("QANYTHING_USER_ID")
        self.headers = qanything_headers
        self.kb_id = kb_id

    def _prepare_request_data(self):
        """准备请求数据"""
        return {"user_id": self.user_id, "kb_id": self.kb_id}

    def _send_request(self, url, data):
        """发送请求"""
        try:
            response = session.post(url, headers=self.headers, data=json.dumps(data))
            return response
        except requests.RequestException as e:
            logger.error(f"请求错误: {e}")
            raise

    def _handle_response(self, response):
        """处理响应"""
        try:
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            logger.error(f"HTTP错误: {http_err}")
        except Exception as err:
            logger.error(f"其他错误: {err}")

    def get_list_files(self):
        """获取指定知识库下所有文件"""
        data = self._prepare_request_data()
        try:
            response = self._send_request(self.url, data)
            return self._handle_response(response)
        except requests.RequestException as e:
            logger.error(f"请求失败: {e}")


class BotInitClient:

    def __init__(
        self,
        kb_ids: list,
        bot_name: str,
        description: str = "",
        prompt_setting: str = conf.get_value("QANYTHING_DEFAULT_PROMPT"),
        welcome_message: str = "您好，我是您的专属机器人，请问有什么可以帮您呢？",
    ) -> None:
        self.url = qanything_api_new_bot
        self.user_id = conf.get_value("QANYTHING_USER_ID")
        self.headers = qanything_headers
        self.kb_ids = kb_ids
        self.bot_name = bot_name
        self.prompt_setting = prompt_setting
        self.description = description
        self.welcome_message = welcome_message

    def _prepare_request_data(self):
        """准备请求数据"""
        return {
            "user_id": self.user_id,
            "bot_name": self.bot_name,
            "description": self.description,
            "prompt_setting": self.prompt_setting,
            "welcome_message": self.welcome_message,
            "model": "MiniChat-2-3B",
            "kb_ids": self.kb_ids,
        }

    def _send_request(self, url, data):
        """发送请求"""
        try:
            response = session.post(url, headers=self.headers, data=json.dumps(data))
            return response
        except requests.RequestException as e:
            logger.error(f"请求错误: {e}")
            raise

    def _handle_response(self, response):
        """处理响应"""
        try:
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            logger.error(f"HTTP错误: {http_err}")
        except Exception as err:
            logger.error(f"其他错误: {err}")

    def init_bot(self):
        """获取指定知识库下所有文件"""
        data = self._prepare_request_data()
        try:
            response = self._send_request(self.url, data)
            return self._handle_response(response)
        except requests.RequestException as e:
            logger.error(f"请求失败: {e}")


class BotInfoClient:

    def __init__(self, bot_id=None) -> None:
        self.url = qanything_api_get_bot_info
        self.user_id = conf.get_value("QANYTHING_USER_ID")
        self.headers = qanything_headers
        self.bot_id = bot_id

    def _prepare_request_data(self):
        """准备请求数据"""

        if self.bot_id:
            return {
                "user_id": self.user_id,
                "bot_id": self.bot_id,
            }
        return {
            "user_id": self.user_id,
        }

    def _send_request(self, url, data):
        """发送请求"""
        try:
            response = session.post(url, headers=self.headers, data=json.dumps(data))
            return response
        except requests.RequestException as e:
            logger.error(f"请求错误: {e}")
            raise

    def _handle_response(self, response):
        """处理响应"""
        try:
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            logger.error(f"HTTP错误: {http_err}")
        except Exception as err:
            logger.error(f"其他错误: {err}")

    def get_bot_info(self):
        """获取指定知识库下所有文件"""
        data = self._prepare_request_data()
        try:
            response = self._send_request(self.url, data)
            return self._handle_response(response)
        except requests.RequestException as e:
            logger.error(f"请求失败: {e}")
            raise
